package com.cyy.tree;

import java.util.Comparator;

/**
 * AVL树，平衡二叉树的一个扩展树
 * @author 会玩的洋洋
 * @since 2022-02-24
 * @param <E>
 */
public class AVLTree<E> extends BBST<E> {

    public AVLTree() {
    }

    public AVLTree(Comparator<E> comparator) {
        super(comparator);
    }

    @Override
    protected void afterAdd(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalance(node)) {
                // 如果添加完是平衡的 则更新高度即可
                updateHeight(node);
            } else {
                // 平衡 AVL 树
                reBalance(node);
                // 只要平衡就跳出循环
                break;
            }
        }
    }

    @Override
    protected void afterRemove(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalance(node)) {
                updateHeight(node);
            } else {
                // 删除无法确保所有祖父节点都平衡
                reBalance(node);
            }
        }
    }

    @Override
    protected Node<E> createNode(E element, Node<E> node) {
        return new AVLNode<>(node, element);
    }

    /**
     * 判断该 AVL 树 是否平衡
     * @param node 需要判断的节点
     * @return boolean
     */
    private boolean isBalance(Node<E> node) {
        return ((AVLNode<E>) node).balanceFactor() <= 1;
    }

    /**
     * 更新节点高度
     * @param node 目标节点
     */
    private void updateHeight(Node<E> node) {
        ((AVLNode<E>) node).updateHeight();
    }

    /**
     * 平衡AVL树
     * @param grand 需要平衡的节点祖父节点
     */
    private void reBalance(Node<E> grand) {
        Node<E> parent = ((AVLNode<E>) grand).tallerNode();
        Node<E> child = ((AVLNode<E>) parent).tallerNode();

        if (parent.isLeftChildren()) {
            // L
            if (child.isRightChildren()) {
                // LR 情况
                leftRotate(parent);
            }
            rightRotate(grand);
        } else {
            // R
            if (child.isLeftChildren()) {
                // RL 情况
                rightRotate(parent);
            }
            leftRotate(grand);
        }
    }

    /**
     * 恢复平衡统一方法
     * @param grand 高度最低的不平衡节点
     */
    protected void reBalance2(Node<E> grand) {
        Node<E> parent = ((AVLTree.AVLNode<E>) grand).tallerNode();
        Node<E> node = ((AVLTree.AVLNode<E>) parent).tallerNode();
        if (parent.isLeftChildren()) {
            if (node.isLeftChildren()) {
                // LL
                rotate(grand, node, node.right, parent, parent.right, grand);
            } else {
                // LR
                rotate(grand, parent, node.left, node, node.right, grand);
            }
        } else {
            if (node.isLeftChildren()) {
                // RL
                rotate(grand, grand, node.left, node, node.right, parent);
            } else {
                // RR
                rotate(grand, grand, parent.left, parent, node.left, node);
            }
        }
    }

    @Override
    protected void rotate(Node<E> r, Node<E> b, Node<E> c, Node<E> d, Node<E> e, Node<E> f) {
        super.rotate(r, b, c, d, e, f);

        updateHeight(b);
        updateHeight(f);
        updateHeight(d);
    }

    /**
     * 更新高度
     * @param grand 祖父节点
     * @param parent 父节点
     * @param child 子节点
     */
    @Override
    protected void afterRotate(Node<E> grand, Node<E> parent, Node<E> child) {
        super.afterRotate(grand, parent, child);

        updateHeight(grand);
        updateHeight(parent);
    }

    private static class AVLNode<E> extends Node<E>{
        int height = 1;

        public AVLNode(Node<E> parent, E element) {
            super(parent, element);
        }

        /**
         * 计算该节点的平衡因子
         * @return int
         */
        public int balanceFactor() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            return Math.abs(leftHeight - rightHeight);
        }

        /**
         * 更新节点高度
         */
        public void updateHeight() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            height = 1 + Math.max(leftHeight, rightHeight);
        }

        /**
         * 获取高度较高的节点
         * @return Node<E>
         */
        public Node<E> tallerNode() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;

            if (leftHeight > rightHeight) {
                return left;
            } else if (leftHeight < rightHeight){
                return right;
            } else {
                // 如果相等， 则判断他是父节点的哪一个子树中的，如果是左子树的 就返回左子节点
                return isLeftChildren() ? left : right;
            }
        }

        @Override
        public String toString() {
            E parentE = null;
            if (parent != null) {
                parentE = parent.element;
            }
            return element + "_p(" + parentE + ")_h(" + height + ")";
        }
    }
}
